/**
 * 
 */
package de.uni_postdam.ling.tcl.smc_lda;

import java.util.ArrayList;
import org.apache.commons.math3.random.RandomGenerator;
import de.uni_postdam.ling.tcl.util.ArraySampler;
import it.unimi.dsi.fastutil.longs.Long2LongOpenHashMap;

/**
 * @author christoph
 *
 */
public class Particle
{
	/**
	 * 
	 */
	private final RandomGenerator rg;
	
	/**
	 * 
	 */
	private final double[] weights;
	
	/**
	 * 
	 */
	private final Long2LongOpenHashMap[] topicWordDistributions;
	
	/**
	 * 
	 */
	private final long[] topicWordSums;
	
	/**
	 * 
	 */
	private final double topicSmoothing;
	
	/**
	 * 
	 */
	private final double topicWordSmoothing;
	
	/**
	 * 
	 */
	private double logImportanceWeight = 0.0;
	
	/**
	 * 
	 */
	private final int gibbsRounds;
	
	/**
	 * 
	 */
	private final DocumentStorage[] history;

	/**
	 * 
	 */
	private final double topicSmoothingSum;
	
	/**
	 * @param topicWordDistributions
	 * @param topicSmoothing
	 * @param topicWordSmoothing
	 */
	public Particle(int numberOfTopics, double topicSmoothing, double topicWordSmoothing,
					int gibbsRounds, int historySize, RandomGenerator rg)
	{
		super();
		this.topicWordDistributions = new Long2LongOpenHashMap[numberOfTopics];
		this.topicWordSums = new long[numberOfTopics];
		
		for(int i=0;i<numberOfTopics;++i)
		{
			this.topicWordDistributions[i] = new Long2LongOpenHashMap();
		}
		
		this.topicSmoothing = topicSmoothing;
		this.topicWordSmoothing = topicWordSmoothing;
		
		this.gibbsRounds = gibbsRounds;
		this.history = new DocumentStorage[historySize];
		
		for(int i=0;i<historySize;++i)
		{
			this.history[i] = new DocumentStorage(numberOfTopics);
		}
		
		this.weights = new double[numberOfTopics];
		this.rg = rg;
		
		this.topicSmoothingSum = topicSmoothing*numberOfTopics;
	}
	
	/**
	 * @return the logImportanceWeight
	 */
	public double getLogImportanceWeight()
	{
		return this.logImportanceWeight;
	}

	/**
	 * 
	 * @param particles
	 * @param cdf
	 * @param samp
	 */
	public void rejuvenate(ArrayList<Particle> particles, double[] cdf, ArraySampler samp)
	{
		Particle parent = particles.get(samp.produceSample(cdf));
		
		for(int i=0;i<this.topicWordDistributions.length;++i)
		{
			Long2LongOpenHashMap map = this.topicWordDistributions[i];
			map.clear();
			map.putAll(parent.topicWordDistributions[i]);
			this.topicWordSums[i] = parent.topicWordSums[i];
		}
		
		for(int i=0;i<this.history.length;++i)
		{
			this.history[i].copy(parent.history[i]);
		}

		boolean seen = false;
		for(int round=0;round < this.gibbsRounds;++round)
		{
			for(int i=0;i<history.length;++i)
			{
				DocumentStorage ds = this.history[i];

				if(ds.getLength() > 0)
				{
					gibbsStep(i, ds);
					seen = true;
				}
			}
			
			if(!seen)
			{
				break;
			}
		}
		
		this.logImportanceWeight = 0.0;
	}
	
	/**
	 * 
	 */
	public void expand(DocumentSource source)
	{
		for(int i=0;i<history.length;++i)
		{
			DocumentStorage ds = this.history[i];
			
			smcStep(source, i, ds);
		}
	}

	/**
	 * 
	 * @param numberOfDocuments
	 * @param numberOfGibbsSteps
	 * @param source
	 */
	public void initialize(int numberOfDocuments, int numberOfGibbsSteps, DocumentSource source)
	{
		DocumentStorage[] stores = new DocumentStorage[numberOfDocuments];
		
		for(int i=0;i<numberOfDocuments;++i)
		{
			stores[i] = new DocumentStorage(this.topicWordSums.length);
			
			this.smcStep(source, i, stores[i]);
			this.logImportanceWeight = 0.0;
		}
		
		for(int step = 0;step < numberOfGibbsSteps;++step)
		for(int i=0;i<numberOfDocuments;++i)
		{
			DocumentStorage ds = stores[i];
			
			this.gibbsStep(i, ds);
		}
	}
	
	
	/**
	 * 
	 * @param source
	 * @param i
	 * @param ds
	 * @return
	 */
	public void smcStep(DocumentSource source, int i, DocumentStorage ds)
	{
		long[] words = source.getDocument(i);
		
		ds.reset(words == null ? new long[0] : words);

		for(int k=0;k<ds.getLength();++k)
		{
			int topic = this.sampleTopic(ds, k, 1.0);
			ds.addTopic(topic);
			ds.addToTopicCount(topic, 1);
			this.addToCount(topic, ds.getWord(k), 1);
		}
	}

	/**
	 * 
	 * @param i
	 * @param ds
	 * @return
	 */
	private void gibbsStep(int i, DocumentStorage ds)
	{
		for(int k=0;k<ds.getLength();++k)
		{
			int topic = ds.getTopic(k);
			ds.addToTopicCount(topic, -1);
			long word = ds.getWord(k);
			this.addToCount(topic, word, -1);
			
			topic = this.sampleTopic(ds, k, 0.0);
			ds.setTopic(k, topic);
			ds.addToTopicCount(topic, 1);
			this.addToCount(topic, ds.getWord(k), 1);
		}
	}

	/**
	 * 
	 * @param topic
	 * @param word
	 * @param amount
	 */
	private void addToCount(int topic, long word, long amount)
	{
		this.topicWordDistributions[topic].addTo(word, amount);
	}

	/**
	 * 
	 * @param ds
	 * @param pos
	 * @param importanceMultiplier
	 * @return
	 */
	private int sampleTopic(DocumentStorage ds, int pos, double importanceMultiplier)
	{
		double max = Double.NEGATIVE_INFINITY;
		
		for(int topic=0;topic<this.topicWordDistributions.length;++topic)
		{
			max = Math.max(max, this.weights[topic] = this.makeWeight(ds,pos,topic));
		}
		
		double sum = 0.0;
		
		for(int topic=0;topic<this.weights.length;++topic)
		{
			sum += this.weights[topic] = Math.exp(weights[topic]-1);
		}
		
		double decision = this.rg.nextDouble();
		
		int choice = -1;
				
		for(;choice<this.weights.length;++choice)
		{
			decision -= this.weights[choice]/sum;
			if(decision <= 0.0)
			{
				break;
			}
		}
		
		if(decision > 0.0)
		{
			throw new IllegalStateException("Probabilities did not sum to one.");
		}
		
		this.logImportanceWeight += importanceMultiplier*(Math.log(sum)-max);
		
		return choice;
	}

	/**
	 * 
	 * @param ds
	 * @param pos
	 * @param topic
	 * @return
	 */
	private double makeWeight(DocumentStorage ds, int pos, int topic)
	{
		long a = this.topicWordDistributions[topic].get(ds.getWord(pos));
		
		double tW = Math.log(a < 1 ? this.topicWordSmoothing : a)
					-Math.log(this.topicWordSmoothing+this.topicWordSums[topic]);
		
		double t = Math.log(ds.getTopicCount(topic)+this.topicSmoothing)
					-Math.log(ds.getSum()+this.topicSmoothingSum);
		
		return t+tW;
	}
	
	/**
	 * 
	 * @return
	 */
	public int getHistorySize()
	{
		return this.history.length;
	}
}
